{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a21bf5f1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/bmp1/anaconda3/envs/mmm/lib/python3.10/site-packages/monai/utils/module.py:396: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.\n",
      "  pkg = __import__(module)  # top level module\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import yaml\n",
    "import os\n",
    "\n",
    "import nibabel as nib\n",
    "import numpy as np\n",
    "import transforms as transforms\n",
    "from torch.utils.data import DataLoader,Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "35259a8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def reader(path: str):\n",
    "    with open(path, 'r', encoding='utf-8') as file:\n",
    "        config = yaml.safe_load(file)\n",
    "    return config\n",
    "\n",
    "def get_data_config(path):\n",
    "    return reader(path)\n",
    "\n",
    "def open_nii(path):\n",
    "    assert path.endswith('.nii') or path.endswith('.nii.gz'),'open_nii function parameter path is must end with .nii or .nii.gz'\n",
    "    data = nib.load(path)\n",
    "    data_array = data.get_fdata()\n",
    "    return data_array\n",
    "\n",
    "class BraTS(torch.utils.data.Dataset):\n",
    "    def __init__(self, config_path='../configs/brats.yaml', is_train=True, resize=[160,176,144]):\n",
    "        super().__init__()\n",
    "        self.config = get_data_config(config_path)['data']\n",
    "        self.is_train = is_train\n",
    "        self.resize = resize\n",
    "        \n",
    "        if self.is_train:\n",
    "            self.data_root = self.config['train_data']\n",
    "        else:\n",
    "            self.data_root = self.config['val_data']\n",
    "        \n",
    "        self.tasks = self.config['tasks']\n",
    "        \n",
    "        self.samples = []\n",
    "        for patient in os.listdir(self.data_root):\n",
    "            patient_path = os.path.join(self.data_root, patient)\n",
    "            if not os.path.isdir(patient_path): \n",
    "                continue\n",
    "                \n",
    "            files = os.listdir(patient_path)\n",
    "            sample_dict = {}\n",
    "            \n",
    "            for file in files:\n",
    "                task_name = file.split('.')[0].split(\"_\")[-1]\n",
    "                if task_name in self.tasks:\n",
    "                    sample_dict[task_name] = os.path.join(patient_path, file)\n",
    "            \n",
    "            if all(task in sample_dict for task in self.tasks):\n",
    "                self.samples.append(sample_dict)\n",
    "        \n",
    "        self.loading_transforms = transforms.Compose([\n",
    "            transforms.LoadImaged(\n",
    "                keys=self.tasks,\n",
    "                image_only=False,\n",
    "                ensure_channel_first=True,\n",
    "                allow_missing_keys=True,\n",
    "            ),\n",
    "            transforms.EnsureTyped(keys=self.tasks, dtype=np.float32),\n",
    "        ])\n",
    "        \n",
    "        self.shape_transform_list = transforms.Compose([\n",
    "            transforms.ScaleIntensityRangePercentilesd(\n",
    "                keys=self.tasks, lower=0.5, upper=99.5, \n",
    "                b_min=-1, b_max=1, clip=True, relative=False, channel_wise=True\n",
    "            ),\n",
    "            transforms.CenterSpatialCropd(\n",
    "                keys=self.tasks, allow_missing_keys=True, roi_size=resize\n",
    "            )\n",
    "        ])\n",
    "\n",
    "        if self.is_train:\n",
    "            self.augmentation_transforms = transforms.Compose([\n",
    "                transforms.RandZoomd(\n",
    "                    keys=self.tasks, allow_missing_keys=True, prob=0.2,\n",
    "                    min_zoom=1.0, max_zoom=1.4,\n",
    "                    mode=['trilinear', 'nearest-exact', 'nearest-exact', 'nearest-exact']\n",
    "                ),\n",
    "                transforms.RandFlipd(\n",
    "                    keys=self.tasks, allow_missing_keys=True, \n",
    "                    prob=0.2, spatial_axis=[0,1]\n",
    "                ),\n",
    "                transforms.RandAdjustContrastd(\n",
    "                    keys=self.tasks, allow_missing_keys=True,\n",
    "                    prob=0.2, gamma=[1.0, 2.0], retain_stats=True\n",
    "                )\n",
    "            ])\n",
    "        else:\n",
    "            self.augmentation_transforms = transforms.Compose([])  \n",
    "        \n",
    "        self.norm_transform_list = transforms.Compose([\n",
    "            transforms.ScaleIntensityRangePercentilesd(\n",
    "                keys=self.tasks, lower=0.5, upper=99.5,\n",
    "                b_min=-1, b_max=1, clip=True, relative=False, channel_wise=True\n",
    "            )\n",
    "        ])\n",
    "\n",
    "        self.total_transforms = transforms.Compose([\n",
    "            self.loading_transforms,\n",
    "            self.shape_transform_list,\n",
    "            self.augmentation_transforms,\n",
    "            self.norm_transform_list,\n",
    "        ])\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.samples)\n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        sample_paths = self.samples[index]\n",
    "        \n",
    "        try:\n",
    "            transformed_data = self.total_transforms(sample_paths)\n",
    "            return transformed_data\n",
    "        except Exception as e:\n",
    "            print(f\"Error transforming sample {index}: {e}\")\n",
    "            return self.__getitem__((index + 1) % len(self.samples))\n",
    "\n",
    "def get_dataloader(data_name, config_path, batch_size=8, shuffle=True, num_workers=8, pin_memory=True, drop_last=True):\n",
    "    assert data_name == 'brats', f'{data_name} is not exist'\n",
    "    \n",
    "    if data_name == 'brats':\n",
    "        train_set = BraTS(config_path,is_train=True)\n",
    "        val_set = BraTS(config_path, is_train=False)\n",
    "    train_loader = DataLoader(train_set,batch_size=batch_size,shuffle=shuffle,num_workers=num_workers,pin_memory=pin_memory,drop_last=drop_last)\n",
    "    val_loader = DataLoader(val_set,batch_size=batch_size,shuffle=False,num_workers=num_workers,pin_memory=pin_memory,drop_last=False)\n",
    "\n",
    "    return train_loader, val_loader\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4e30ed13",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/bmp1/anaconda3/envs/mmm/lib/python3.10/site-packages/torch/utils/data/dataloader.py:557: UserWarning: This DataLoader will create 8 worker processes in total. Our suggested max number of worker in current system is 4, which is smaller than what this DataLoader is going to create. Please be aware that excessive worker creation might get DataLoader running slow or even freeze, lower the worker number to avoid potential slowness/freeze if necessary.\n",
      "  warnings.warn(_create_warning_msg(\n"
     ]
    }
   ],
   "source": [
    "train_loader, val_loader = get_dataloader('brats', '../configs/brats.yaml')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9e5032d0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/bmp1/anaconda3/envs/mmm/lib/python3.10/site-packages/torch/utils/data/dataloader.py:557: UserWarning: This DataLoader will create 8 worker processes in total. Our suggested max number of worker in current system is 4, which is smaller than what this DataLoader is going to create. Please be aware that excessive worker creation might get DataLoader running slow or even freeze, lower the worker number to avoid potential slowness/freeze if necessary.\n",
      "  warnings.warn(_create_warning_msg(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([8, 1, 160, 176, 144])\n"
     ]
    }
   ],
   "source": [
    "for data in train_loader:\n",
    "    print(data['t1'].shape)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d141d539",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mmm",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
